package com.google.code.shardbatis.plugin;

import java.io.IOException;
import java.io.InputStream;
import java.lang.reflect.Method;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

import org.apache.ibatis.executor.statement.RoutingStatementHandler;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.io.Resources;
import org.apache.ibatis.logging.Log;
import org.apache.ibatis.logging.LogFactory;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.ResultMap;
import org.apache.ibatis.mapping.ResultMapping;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Plugin;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.session.ResultHandler;

import com.google.code.shardbatis.builder.ShardConfigHolder;
import com.google.code.shardbatis.builder.ShardConfigParser;
import com.google.code.shardbatis.converter.SqlConverterFactory;
import com.google.code.shardbatis.converter.SqlConverterOrderBy;
import com.google.code.shardbatis.util.CustomComparator;
import com.google.code.shardbatis.util.FourTuples;
import com.google.code.shardbatis.util.ReflectionUtils;

@Intercepts({ 
	@Signature(
		type = StatementHandler.class, 
		method = "update", 
		args = { java.sql.Statement.class }
	)
	,
	@Signature(
		type = StatementHandler.class, 
		method = "query", 
		args = { java.sql.Statement.class, ResultHandler.class }
	)
})
public class ShardPlugin implements Interceptor {
	
	private static final Log log = LogFactory.getLog(ShardPlugin.class);
	
	private static final int QUERY_MAX_RESULT = 500;
	
	public static final String SHARDING_CONFIG = "shardingConfig";
	private static final ConcurrentHashMap<String, Boolean> cache = new ConcurrentHashMap<String, Boolean>();

	@SuppressWarnings({ "rawtypes" })
	public Object intercept(Invocation invocation) throws Throwable {
		Object target = invocation.getTarget();
		Object[] args = invocation.getArgs();
		
		if (target instanceof StatementHandler) {
			StatementHandler statementHandler = (StatementHandler)target;
			
			Method method = invocation.getMethod();
			
			MappedStatement mappedStatement = null;
			if (statementHandler instanceof RoutingStatementHandler) {
				StatementHandler delegate = (StatementHandler) ReflectionUtils.getFieldValue(statementHandler, "delegate");
				mappedStatement = (MappedStatement) ReflectionUtils.getFieldValue(delegate, "mappedStatement");
			}
			else {
				mappedStatement = (MappedStatement) ReflectionUtils.getFieldValue(statementHandler, "mappedStatement");
			}
			
			String mapperId = mappedStatement.getId();
			//判断是否于需要进行分表处理
			//若不需要，按源流程记性；若需要，则按分表模式进行
			if (isShouldParse(mapperId)) {
				
				//获取原始SQL
				String osql = statementHandler.getBoundSql().getSql();
				if (log.isDebugEnabled()) {
					log.debug("Original Sql [" + mapperId + "]:" + osql);
				}
				
				//获取当前SQL声明
				Statement stmt = (Statement)invocation.getArgs()[0];
				//获取SQL参数
				Object params = statementHandler.getBoundSql().getParameterObject();
				
				//SQL转换
				SqlConverterFactory cf = SqlConverterFactory.getInstance();
				FourTuples<List<SqlConverterOrderBy>, String[], String, String[]> ft = cf.convert(osql, params, mapperId);
				
				Set<String> availableTables = getAvailableTables(stmt, ft.c);
				
				//执行更新（增、删、改）
				if (method.getName().equals("update")) {
					int count = 0;
					for (int i = 0; i < ft.b.length; i++) {
						if (null == ft.b[i]) { continue; }
						if (null == ft.d[i]) { continue; }
						if (!availableTables.contains(ft.d[i])) { continue; }
						if (log.isDebugEnabled()) {
							log.debug("Converted Update Sql [" + mapperId + "]:" + ft.b[i]);
						}
						ReflectionUtils.setFieldValue(statementHandler.getBoundSql(), "sql", ft.b[i]);
						Statement newStmt = statementHandler.prepare(stmt.getConnection());
						statementHandler.parameterize(newStmt);
						count += statementHandler.update(newStmt);
					}
					return count;
				}
				//执行查询
				else if (method.getName().equals("query")) {
					List<Object> list = new ArrayList<Object>();
					loop1:for (int i = 0; i < ft.b.length; i++) {
						if (null == ft.b[i]) { continue; }
						if (null == ft.d[i]) { continue; }
						if (!availableTables.contains(ft.d[i])) { continue; }
						if (log.isDebugEnabled()) {
							log.debug("Converted Query Sql [" + mapperId + "]:" + ft.b[i]);
						}
						ReflectionUtils.setFieldValue(statementHandler.getBoundSql(), "sql", ft.b[i]);
						Statement newStmt = statementHandler.prepare(stmt.getConnection());
						statementHandler.parameterize(newStmt);
						List<Object> objs = statementHandler.query(newStmt, (ResultHandler)args[1]);
						if (null != objs) {
							//仅返回QUERY_MAX_RESULT条记录
							//若已查询记录数 > QUERY_MAX_RESULT,则返回
							//反之，继续循环查询
							for (Object obj : objs) {
								if (list.size() > QUERY_MAX_RESULT) {
									break loop1;
								}
								list.add(obj);
							}
						}
					}
					if (null != ft.a && ft.a.size() > 0 && null != mappedStatement) {
						List<ResultMap> resultMaps = mappedStatement.getResultMaps();
						if (null != resultMaps && resultMaps.size() > 0) {
							parse(ft.a, resultMaps.get(0).getResultMappings());
							Collections.sort(list, new CustomComparator(ft.a));
						}
					}
					if (mapperId.toLowerCase().endsWith("_count")) {
						Long count = 0L;
						for (Object c : list) {
							count += Long.parseLong(c.toString());
						}
						list = new ArrayList<Object>();
						list.add(count);
					} 
					return list;
				}
			}
		}
		return invocation.proceed();
	}
	
	private void parse(List<SqlConverterOrderBy> orderBys, List<ResultMapping> resultMappings) {
		Map<String, String> mappings = new HashMap<String, String>();
		for (ResultMapping resultMapping : resultMappings) {
			mappings.put(resultMapping.getColumn(), resultMapping.getProperty());
		}
		for (SqlConverterOrderBy orderBy : orderBys) {
			if (mappings.containsKey(orderBy.getColumnName())) {
				orderBy.setFieldName(mappings.get(orderBy.getColumnName()));
			}
		}
	}
	
	public Object plugin(Object target) {
		return Plugin.wrap(target, this);
	}

	public void setProperties(Properties properties) {
		String config = properties.getProperty("shardingConfig", null);
		log.debug("shardingConfig = " + config);
		if ((config == null) || (config.trim().length() == 0)) {
			throw new IllegalArgumentException("property 'shardingConfig' is requested.");
		}

		ShardConfigParser parser = new ShardConfigParser();
		InputStream input = null;
		try {
			input = Resources.getResourceAsStream(config);
			parser.parse(input);
		} catch (IOException e) {
		} catch (Exception e) {
		} finally {
			if (input != null)
				try {
					input.close();
				} catch (IOException e) {
					log.error(e.getMessage(), e);
				}
		}
	}
	
	private Set<String> getAvailableTables(Statement stmt, String baseTableName) {
		Set<String> s = new HashSet<String>();
		Connection conn = null;
		ResultSet rs = null;
		try {
			conn = stmt.getConnection();
			DatabaseMetaData metadata = conn.getMetaData();
			rs = metadata.getTables(null, null, baseTableName + "%", new String[]{"TABLE"});
			while (rs.next()) {
				s.add(rs.getString("TABLE_NAME"));
			}
		}
		catch (Throwable e) {
			log.error("getAvailableTables err", e);
		}
		finally {
			if (null != rs) {
				try { rs.close(); } 
				catch (Exception e) { }
			}
		}
		return s;
	} 

	private boolean isShouldParse(String mapperId) {
		Boolean parse = (Boolean) cache.get(mapperId);

		if (parse != null) {
			return parse.booleanValue();
		}

		if (!(mapperId.endsWith("!selectKey"))) {
			ShardConfigHolder configHolder = ShardConfigHolder.getInstance();

			if ((!configHolder.isIgnoreId(mapperId)) && configHolder.isParseId(mapperId)) {
				parse = Boolean.valueOf(true);
			}
		}

		if (parse == null) {
			parse = Boolean.valueOf(false);
		}
		cache.put(mapperId, parse);
		return parse.booleanValue();
	}
}